#!/bin/bash

# ThingsBoard 还原脚本
# 作者：Vondser
# 版本：1.0
# 功能：还原ThingsBoard的PostgreSQL数据库、Cassandra数据库、配置文件和数据文件

# 日志配置
LOG_DIR="/var/log/tb-backup"
LOG_FILE="${LOG_DIR}/tb_restore_$(date +"%Y-%m-%d").log"
SCRIPT_NAME=$(basename "$0")

# 确保日志目录存在
mkdir -p ${LOG_DIR}

# 日志函数
log_info() {
    local message="[INFO] [$(date +"%Y-%m-%d %H:%M:%S")] [${SCRIPT_NAME}] $1"
    echo -e "\033[0;32m${message}\033[0m" | tee -a ${LOG_FILE}
}

log_warning() {
    local message="[WARNING] [$(date +"%Y-%m-%d %H:%M:%S")] [${SCRIPT_NAME}] $1"
    echo -e "\033[0;33m${message}\033[0m" | tee -a ${LOG_FILE}
}

log_error() {
    local message="[ERROR] [$(date +"%Y-%m-%d %H:%M:%S")] [${SCRIPT_NAME}] $1"
    echo -e "\033[0;31m${message}\033[0m" | tee -a ${LOG_FILE}
}

log_success() {
    local message="[SUCCESS] [$(date +"%Y-%m-%d %H:%M:%S")] [${SCRIPT_NAME}] $1"
    echo -e "\033[0;32m${message}\033[0m" | tee -a ${LOG_FILE}
}

# 检查是否以root用户运行
if [ "$(id -u)" != "0" ]; then
   log_error "此脚本必须以root用户身份运行"
   exit 1
fi

# 检查参数
if [ $# -ne 1 ]; then
    log_error "用法: $0 <备份文件路径>"
    echo "用法: $0 <备份文件路径>"
    echo "示例: $0 /backup/thingsboard/thingsboard_backup_2023-04-01_02-00-00.tar.gz"
    exit 1
fi

BACKUP_FILE=$1

# 记录脚本开始执行
log_info "====== ThingsBoard还原开始 ======"
log_info "备份文件: ${BACKUP_FILE}"
log_info "日志文件: ${LOG_FILE}"

# 检查备份文件是否存在
if [ ! -f "${BACKUP_FILE}" ]; then
    log_error "备份文件不存在: ${BACKUP_FILE}"
    exit 1
fi

# 检查备份文件的MD5校验和
if [ -f "${BACKUP_FILE}.md5" ]; then
    log_info "验证备份文件完整性..."
    EXPECTED_MD5=$(cat ${BACKUP_FILE}.md5 | awk '{print $1}')
    ACTUAL_MD5=$(md5sum ${BACKUP_FILE} | awk '{print $1}')
    
    if [ "$EXPECTED_MD5" = "$ACTUAL_MD5" ]; then
        log_success "备份文件MD5校验通过"
    else
        log_error "备份文件MD5校验失败！预期：${EXPECTED_MD5}，实际：${ACTUAL_MD5}"
        read -p "文件可能已损坏，是否继续？(y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_error "用户取消还原"
            exit 1
        fi
        log_warning "用户选择继续还原，但文件可能已损坏"
    fi
else
    log_warning "未找到MD5校验文件，无法验证备份文件完整性"
    read -p "是否继续？(y/n): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_error "用户取消还原"
        exit 1
    fi
    log_warning "用户选择继续还原"
fi

# 创建临时目录
log_info "创建临时目录..."
TEMP_DIR=$(mktemp -d)
log_info "临时目录: ${TEMP_DIR}"

# 解压备份文件
log_info "解压备份文件..."
tar -xzf ${BACKUP_FILE} -C ${TEMP_DIR}
if [ $? -ne 0 ]; then
    log_error "备份文件解压失败"
    rm -rf ${TEMP_DIR}
    exit 1
fi
log_success "备份文件解压成功"

# 获取备份日期
BACKUP_DATE=$(basename ${BACKUP_FILE} | sed 's/thingsboard_backup_//' | sed 's/.tar.gz//')
EXTRACT_DIR="${TEMP_DIR}/${BACKUP_DATE}"

log_info "备份日期: ${BACKUP_DATE}"
log_info "解压目录: ${EXTRACT_DIR}"

# 检查解压后的目录结构
log_info "检查备份文件内容..."
required_files=("tb_postgres.sql" "tb_config.tar.gz" "tb_data.tar.gz")
missing_files=""

for file in "${required_files[@]}"; do
    if [ ! -f "${EXTRACT_DIR}/${file}" ]; then
        missing_files="${missing_files} ${file}"
    fi
done

if [ ! -d "${EXTRACT_DIR}/cassandra" ]; then
    missing_files="${missing_files} cassandra/"
fi

if [ -n "$missing_files" ]; then
    log_error "备份文件不完整，缺少以下文件:${missing_files}"
    read -p "是否继续还原？(y/n): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_error "用户取消还原"
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    log_warning "用户选择继续还原，但备份文件不完整"
fi

# 记录备份文件的基本信息
if [ -f "${EXTRACT_DIR}/system_info.txt" ]; then
    log_info "备份系统信息:"
    cat "${EXTRACT_DIR}/system_info.txt" | tee -a ${LOG_FILE}
fi

if [ -f "${EXTRACT_DIR}/package_versions.txt" ]; then
    log_info "备份的软件包版本:"
    cat "${EXTRACT_DIR}/package_versions.txt" | tee -a ${LOG_FILE}
fi

# 确认还原操作
log_warning "还原操作将覆盖当前ThingsBoard数据！"
read -p "是否继续还原操作？(y/n): " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    log_error "用户取消还原"
    rm -rf ${TEMP_DIR}
    exit 1
fi
log_info "用户确认继续还原操作"

# 1. 停止ThingsBoard服务
log_info "停止ThingsBoard服务..."
systemctl stop thingsboard
if [ $? -eq 0 ]; then 
    log_success "ThingsBoard服务已停止"
else 
    log_error "无法停止ThingsBoard服务，请检查服务状态"
    read -p "是否继续还原操作？(y/n): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_error "用户取消还原"
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    log_warning "用户选择继续还原"
fi

# 2. 还原PostgreSQL数据库
if [ -f "${EXTRACT_DIR}/tb_postgres.sql" ]; then
    log_info "还原PostgreSQL数据库..."
    
    # 检查PostgreSQL服务状态
    if systemctl is-active --quiet postgresql; then
        log_info "PostgreSQL服务正在运行"
    else
        log_info "启动PostgreSQL服务..."
        systemctl start postgresql
        if [ $? -ne 0 ]; then
            log_error "无法启动PostgreSQL服务"
            read -p "是否继续还原操作？(y/n): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                log_error "用户取消还原"
                rm -rf ${TEMP_DIR}
                exit 1
            fi
            log_warning "用户选择继续还原，但可能会失败"
        fi
    fi
    
    # 首先删除现有数据库
    log_info "删除现有thingsboard数据库..."
    sudo -u postgres psql -c "DROP DATABASE IF EXISTS thingsboard;" 2>>${LOG_FILE}
    if [ $? -ne 0 ]; then
        log_error "无法删除现有thingsboard数据库"
        log_info "尝试断开现有连接后再删除..."
        sudo -u postgres psql -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname='thingsboard';" 2>>${LOG_FILE}
        sudo -u postgres psql -c "DROP DATABASE IF EXISTS thingsboard;" 2>>${LOG_FILE}
        if [ $? -ne 0 ]; then
            log_error "仍然无法删除thingsboard数据库，还原失败"
            rm -rf ${TEMP_DIR}
            exit 1
        fi
    fi
    
    log_info "创建新的thingsboard数据库..."
    sudo -u postgres psql -c "CREATE DATABASE thingsboard;" 2>>${LOG_FILE}
    if [ $? -ne 0 ]; then
        log_error "无法创建新的thingsboard数据库"
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    
    # 还原数据
    log_info "导入PostgreSQL数据..."
    sudo -u postgres psql thingsboard < ${EXTRACT_DIR}/tb_postgres.sql 2>>${LOG_FILE}
    if [ $? -eq 0 ]; then 
        log_success "PostgreSQL还原成功"
    else 
        log_error "PostgreSQL还原失败"
        read -p "是否继续还原其他组件？(y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_error "用户取消还原"
            rm -rf ${TEMP_DIR}
            exit 1
        fi
        log_warning "用户选择继续还原其他组件"
    fi
else
    log_warning "备份中缺少PostgreSQL数据文件，跳过PostgreSQL还原"
fi

# 3. 停止Cassandra服务
log_info "停止Cassandra服务..."
systemctl stop cassandra
if [ $? -eq 0 ]; then 
    log_success "Cassandra服务已停止"
else 
    log_error "无法停止Cassandra服务，请检查服务状态"
    read -p "是否继续还原操作？(y/n): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_error "用户取消还原"
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    log_warning "用户选择继续还原"
fi

# 4. 还原Cassandra数据
if [ -d "${EXTRACT_DIR}/cassandra" ]; then
    log_info "还原Cassandra数据..."
    
    # 4.1 查找Cassandra数据目录
    log_info "查找Cassandra数据目录..."
    CASSANDRA_DATA_DIR=$(grep -r "data_file_directories" /etc/cassandra/cassandra.yaml | awk '{print $2}' | tr -d '[]" ')
    if [ -z "$CASSANDRA_DATA_DIR" ]; then
        log_warning "无法从配置文件获取Cassandra数据目录，使用默认路径"
        CASSANDRA_DATA_DIR="/var/lib/cassandra/data"
    fi
    log_info "Cassandra数据目录: ${CASSANDRA_DATA_DIR}"
    
    # 4.2 备份现有Cassandra数据
    if [ -d "${CASSANDRA_DATA_DIR}/thingsboard" ]; then
        log_info "备份当前Cassandra数据..."
        BAK_DIR="${CASSANDRA_DATA_DIR}/thingsboard.bak.$(date +%Y%m%d%H%M%S)"
        mv ${CASSANDRA_DATA_DIR}/thingsboard ${BAK_DIR}
        if [ $? -eq 0 ]; then
            log_success "现有Cassandra数据已备份到: ${BAK_DIR}"
        else
            log_error "无法备份现有Cassandra数据"
            read -p "是否继续还原操作？(y/n): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                log_error "用户取消还原"
                rm -rf ${TEMP_DIR}
                exit 1
            fi
            log_warning "用户选择继续还原"
        fi
    else
        log_warning "没有找到现有的Cassandra thingsboard数据"
    fi
    
    # 4.3 创建新目录结构
    log_info "创建新的Cassandra数据目录..."
    mkdir -p ${CASSANDRA_DATA_DIR}/thingsboard
    if [ $? -ne 0 ]; then
        log_error "无法创建Cassandra数据目录"
        # 尝试恢复之前的备份
        if [ -d "${BAK_DIR}" ]; then
            log_info "恢复之前的Cassandra数据备份..."
            mv ${BAK_DIR} ${CASSANDRA_DATA_DIR}/thingsboard
        fi
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    
    # 4.4 还原Cassandra数据
    log_info "复制Cassandra备份数据..."
    cp -r ${EXTRACT_DIR}/cassandra/* ${CASSANDRA_DATA_DIR}/thingsboard/ 2>>${LOG_FILE}
    if [ $? -eq 0 ]; then 
        log_success "Cassandra数据复制成功"
    else 
        log_error "Cassandra数据复制失败"
        # 尝试恢复之前的备份
        if [ -d "${BAK_DIR}" ]; then
            log_info "恢复之前的Cassandra数据备份..."
            rm -rf ${CASSANDRA_DATA_DIR}/thingsboard
            mv ${BAK_DIR} ${CASSANDRA_DATA_DIR}/thingsboard
        fi
        read -p "是否继续还原其他组件？(y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_error "用户取消还原"
            rm -rf ${TEMP_DIR}
            exit 1
        fi
        log_warning "用户选择继续还原其他组件"
    fi
    
    # 4.5 设置正确的所有权
    log_info "修复Cassandra数据权限..."
    chown -R cassandra:cassandra ${CASSANDRA_DATA_DIR}/thingsboard
    if [ $? -eq 0 ]; then 
        log_success "Cassandra数据权限修复成功"
    else 
        log_error "Cassandra数据权限修复失败，可能会导致服务无法启动"
    fi
else
    log_warning "备份中缺少Cassandra数据目录，跳过Cassandra还原"
fi

# 5. 还原ThingsBoard配置
if [ -f "${EXTRACT_DIR}/tb_config.tar.gz" ]; then
    log_info "还原ThingsBoard配置..."
    
    # 备份当前配置
    if [ -d "/etc/thingsboard" ]; then
        log_info "备份当前ThingsBoard配置..."
        CONFIG_BAK_DIR="/etc/thingsboard.bak.$(date +%Y%m%d%H%M%S)"
        cp -r /etc/thingsboard ${CONFIG_BAK_DIR}
        if [ $? -eq 0 ]; then
            log_success "当前ThingsBoard配置已备份到: ${CONFIG_BAK_DIR}"
        else
            log_warning "无法备份当前ThingsBoard配置"
        fi
    fi
    
    # 解压配置文件
    tar -xzf ${EXTRACT_DIR}/tb_config.tar.gz -C / 2>>${LOG_FILE}
    if [ $? -eq 0 ]; then 
        log_success "ThingsBoard配置还原成功"
    else 
        log_error "ThingsBoard配置还原失败"
        # 尝试恢复之前的备份
        if [ -d "${CONFIG_BAK_DIR}" ]; then
            log_info "恢复之前的ThingsBoard配置备份..."
            rm -rf /etc/thingsboard
            cp -r ${CONFIG_BAK_DIR} /etc/thingsboard
        fi
        read -p "是否继续还原其他组件？(y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_error "用户取消还原"
            rm -rf ${TEMP_DIR}
            exit 1
        fi
        log_warning "用户选择继续还原其他组件"
    fi
else
    log_warning "备份中缺少ThingsBoard配置文件，跳过配置还原"
fi

# 6. 还原ThingsBoard数据
if [ -f "${EXTRACT_DIR}/tb_data.tar.gz" ]; then
    log_info "还原ThingsBoard数据..."
    
    # 备份当前数据
    if [ -d "/var/lib/thingsboard" ]; then
        log_info "备份当前ThingsBoard数据..."
        DATA_BAK_DIR="/var/lib/thingsboard.bak.$(date +%Y%m%d%H%M%S)"
        cp -r /var/lib/thingsboard ${DATA_BAK_DIR}
        if [ $? -eq 0 ]; then
            log_success "当前ThingsBoard数据已备份到: ${DATA_BAK_DIR}"
        else
            log_warning "无法备份当前ThingsBoard数据"
        fi
    fi
    
    # 解压数据文件
    tar -xzf ${EXTRACT_DIR}/tb_data.tar.gz -C / 2>>${LOG_FILE}
    if [ $? -eq 0 ]; then 
        log_success "ThingsBoard数据还原成功"
    else 
        log_error "ThingsBoard数据还原失败"
        # 尝试恢复之前的备份
        if [ -d "${DATA_BAK_DIR}" ]; then
            log_info "恢复之前的ThingsBoard数据备份..."
            rm -rf /var/lib/thingsboard
            cp -r ${DATA_BAK_DIR} /var/lib/thingsboard
        fi
        read -p "是否继续？(y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_error "用户取消还原"
            rm -rf ${TEMP_DIR}
            exit 1
        fi
        log_warning "用户选择继续"
    fi
else
    log_warning "备份中缺少ThingsBoard数据文件，跳过数据还原"
fi

# 7. 修复权限
log_info "修复文件权限..."
if [ -d "/etc/thingsboard" ]; then
    chown -R thingsboard:thingsboard /etc/thingsboard
    chmod -R 755 /etc/thingsboard
    log_success "ThingsBoard配置权限已修复"
else
    log_warning "ThingsBoard配置目录不存在，跳过权限修复"
fi

if [ -d "/var/lib/thingsboard" ]; then
    chown -R thingsboard:thingsboard /var/lib/thingsboard
    chmod -R 755 /var/lib/thingsboard
    log_success "ThingsBoard数据权限已修复"
else
    log_warning "ThingsBoard数据目录不存在，跳过权限修复"
fi

# 8. 启动服务
log_info "启动Cassandra服务..."
systemctl start cassandra
if [ $? -eq 0 ]; then 
    log_success "Cassandra服务已启动"
else 
    log_error "无法启动Cassandra服务，请检查服务状态"
    read -p "是否继续启动ThingsBoard服务？(y/n): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_error "用户取消启动ThingsBoard服务"
        rm -rf ${TEMP_DIR}
        exit 1
    fi
    log_warning "用户选择继续启动ThingsBoard服务，但可能会失败"
fi

# 等待Cassandra完全启动
log_info "等待Cassandra服务完全启动..."
for i in {1..6}; do
    log_info "等待中... $i/6"
    sleep 10
    
    # 检查Cassandra是否响应
    nodetool status > /dev/null 2>&1
    if [ $? -eq 0 ]; then
        log_success "Cassandra服务已就绪"
        break
    fi
    
    if [ $i -eq 6 ]; then
        log_warning "等待Cassandra超时，但仍将继续"
    fi
done

log_info "启动ThingsBoard服务..."
systemctl start thingsboard
if [ $? -eq 0 ]; then 
    log_success "ThingsBoard服务已启动"
else 
    log_error "无法启动ThingsBoard服务，请查看系统日志以获取详细信息"
    systemctl status thingsboard -l | tee -a ${LOG_FILE}
fi

# 9. 清理
log_info "清理临时文件..."
rm -rf ${TEMP_DIR}
log_success "临时文件清理完成"

# 提示保留的备份
if [ -d "${BAK_DIR}" ]; then
    log_info "Cassandra旧数据备份保留在: ${BAK_DIR}"
    log_info "确认系统正常运行后可手动删除"
fi

if [ -d "${CONFIG_BAK_DIR}" ]; then
    log_info "ThingsBoard配置备份保留在: ${CONFIG_BAK_DIR}"
    log_info "确认系统正常运行后可手动删除"
fi

if [ -d "${DATA_BAK_DIR}" ]; then
    log_info "ThingsBoard数据备份保留在: ${DATA_BAK_DIR}"
    log_info "确认系统正常运行后可手动删除"
fi

# 10. 还原完成，统计信息
RESTORE_END_TIME=$(date +"%Y-%m-%d %H:%M:%S")
RESTORE_DURATION=$(($(date +%s) - $(date -d "$(head -1 ${LOG_FILE} | grep -o '\[INFO\] \[[^]]*\]' | cut -d'[' -f3 | tr -d ']')" +%s)))
DURATION_MINS=$((RESTORE_DURATION / 60))
DURATION_SECS=$((RESTORE_DURATION % 60))

log_info "====== ThingsBoard还原完成 ======"
log_info "还原开始时间: $(head -1 ${LOG_FILE} | grep -o '\[INFO\] \[[^]]*\]' | cut -d'[' -f3 | tr -d ']')"
log_info "还原结束时间: ${RESTORE_END_TIME}"
log_info "还原总耗时: ${DURATION_MINS}分${DURATION_SECS}秒"
log_info "请验证ThingsBoard系统是否正常运行"

# 提示验证步骤
log_info "验证步骤:"
log_info "1. 打开ThingsBoard Web界面"
log_info "2. 使用管理员账户登录"
log_info "3. 检查设备、仪表板、规则链等是否正常"
log_info "4. 检查遥测数据是否可用"

exit 0 